from sqlalchemy.orm import Session
from typing import Any, Dict, Union

from app.models.generation_job import GenerationJob
from app.schemas.generation_job import GenerationJobCreate, GenerationJobUpdate

def get_job(db: Session, *, job_id: int) -> GenerationJob | None:
    return db.query(GenerationJob).filter(GenerationJob.id == job_id).first()

def get_public_jobs(db: Session, skip: int = 0, limit: int = 100) -> list[GenerationJob]:
    return (
        db.query(GenerationJob)
        .filter(GenerationJob.is_public == True)
        .order_by(GenerationJob.created_at.desc())
        .offset(skip)
        .limit(limit)
        .all()
    )

def get_user_jobs(db: Session, *, user_id: int, skip: int = 0, limit: int = 100) -> list[GenerationJob]:
    return (
        db.query(GenerationJob)
        .filter(GenerationJob.user_id == user_id)
        .order_by(GenerationJob.created_at.desc())
        .offset(skip)
        .limit(limit)
        .all()
    )

def create_generation_job(db: Session, *, user_id: int, job_in: GenerationJobCreate) -> GenerationJob:
    db_obj = GenerationJob(
        **job_in.model_dump(),
        user_id=user_id
    )
    db.add(db_obj)
    db.commit()
    db.refresh(db_obj)
    return db_obj

def update_job(
    db: Session, *, db_obj: GenerationJob, obj_in: Union[GenerationJobUpdate, Dict[str, Any]]
) -> GenerationJob:
    if isinstance(obj_in, dict):
        update_data = obj_in
    else:
        update_data = obj_in.model_dump(exclude_unset=True)
    
    for field in update_data:
        if hasattr(db_obj, field):
            setattr(db_obj, field, update_data[field])
            
    db.add(db_obj)
    db.commit()
    db.refresh(db_obj)
    return db_obj
